package cn.itcast.manus.agent;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.resource.ResourceUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import cn.itcast.manus.agent.bean.ReActInput;
import cn.itcast.manus.config.ModelConfig;
import cn.itcast.manus.config.ReActConfig;
import cn.itcast.manus.message.MessageSession;
import cn.itcast.manus.util.JsonFinder;
import jakarta.annotation.Resource;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.DefaultToolDefinition;
import org.springframework.ai.tool.definition.ToolDefinition;

import javax.validation.constraints.NotNull;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.function.Function;

@Slf4j
public abstract class ReActBaseAgent extends BaseAgent {
    public static final String AGENT_OUTPUT_METHOD = "AgentOutput";
    @Resource
    protected ReActConfig reActConfig;
    @Resource(name = ModelConfig.MAIN_AGENT)
    private ChatModel chatModel;
    @Resource(name = ModelConfig.MAIN_AGENT_MODEL_CONFIG)
    private ModelConfig.BaseModelConfig modelConfig;
    @Getter
    protected int currentStep = 1;
    // 是否停止标识
    private boolean agentStop = false;
    // 最终结果
    private String finalResult;

    protected MessageSession messageSession;
    // 工具定义结构内容
    private String toolSchema;
    // 工具定义
    private ToolDefinition toolDefinition;
    // 定义动作映射，key为动作名称，value为动作执行函数（key：参数，value：执行返回值）
    private final Map<String, Function<Map<String, Object>, String>> actionRegistry = new HashMap<>();
    public ReActBaseAgent(MessageSession messageSession) {
        this(messageSession, "schema/schemaBaseReAct.json");
    }
    public ReActBaseAgent(MessageSession messageSession, String toolSchemaPath) {
        this.messageSession = messageSession;
        this.toolSchema = ResourceUtil.readStr(toolSchemaPath, StandardCharsets.UTF_8);
        this.refreshToolDefinition();
    }

    @Override
    protected String solve(String task) {
        return this.reActSolve(task);
    }

    protected String reActSolve(String task) {
        // 初始化当前步骤
        this.currentStep = 1;
        // 初始化动作列表
        this.initActionMap();
        // 定义输入消息列表
        List<Message> inputList = new LinkedList<>();
        // 将当前任务加入到列表中
        this.addInitInput(inputList, task);
        // 定义 token 计数器
        var tokenCount = new AtomicLong(0L);
        // 如果没有标记为停止, 并且当前的步骤小于等于最大步骤，则进行循环
        while (!this.agentStop && this.currentStep <= this.reActConfig.getMaxStep()) {
            this.buildCurrentParse(inputList); //一个消息列表（List<Message>），包含了要发送给大模型的对话历史和指令
            this.callingLLM(inputList, tokenCount);
            log.info("step:{},tokenCount:{}", this.currentStep, tokenCount.get());
            log.info("[{}👆]------------------", getClass().getSimpleName());
            this.currentStep++;
        }
        // 最后一步检查
        if (this.currentStep >= this.reActConfig.getMaxStep() && !this.agentStop) {
            String text = """
                    注意，你的执行已达到最大步数限制，请立即使用done操作返回截至到目前的工作成果，
                    需要包含具体的内容，不允许任何省略，最后总结尚未完成的任务并对遇到的问题加以说明
                    """;
            inputList.add(new UserMessage(text));
            this.callingLLM(inputList, tokenCount);
        }
        return finalResult;
    }
    protected void refreshToolDefinition(){
        JSONObject jsonObject = JSONUtil.parseObj(this.toolSchema);
        this.toolDefinition = DefaultToolDefinition.builder()
                .name(jsonObject.getStr("name"))
                .description(jsonObject.getStr("description"))
                .inputSchema(jsonObject.getStr("inputSchema"))
                .build();
    }
    /**
     * 调用 LLM
     */
    private void callingLLM(List<Message> inputList, AtomicLong tokenCount) {
        ChatResponse chatResponse;
        try {
            // 1. 构造请求参数
            var chatOptions = OpenAiChatOptions.builder()
                    .toolCallbacks(this.getToolCallbacks()) // 设置工具回调
                    .internalToolExecutionEnabled(false) //关闭springai执行tool，由我们自己的程序来执行action
                    .toolChoice(OpenAiApi.ChatCompletionRequest.ToolChoiceBuilder.FUNCTION(AGENT_OUTPUT_METHOD)) //强制设置工具的名字
                    .build();

            // 2. 构造prompt
            var prompt = new Prompt(inputList, chatOptions); //大模型调用的桥梁：在 callingLLM 方法中，这个 Prompt 对象被传递给 chatModel() 方法，作为调用 OpenAI 等大语言模型的请求载体

            // 3. 调用大模型
            if (this.modelConfig.isStream()) {
                //可以简化为
                //chatResponse = this.chatModel().stream(prompt).blockFirst();
                // 流式调用
                chatResponse = this.chatModel()
                        .stream(prompt)
                        .collectList()//使用 Reactor 的 collectList() 操作符将流中的所有 ChatResponse 对象收集到一个列表中
                        .block()//调用 block() 方法阻塞当前线程，直到异步操作完成并获取结果 返回收集到的 List<ChatResponse> 对象
                        .stream()//List<ChatResponse> 转换为 Java 8 的 Stream<ChatResponse> 流
                        .findFirst()//获取流中的第一个 ChatResponse 对象返回 Optional<ChatResponse> 类型
                        .orElseThrow(); //如果 Optional 为空（没有响应），则抛出异常否则返回 ChatResponse 对象
            } else {
                // 非流式
                chatResponse = this.chatModel().call(prompt);
            }
        } catch (Exception e) {
            log.error("调用AI大模型时出错！", e);
            throw new RuntimeException(e);
        }

        // 4. 处理大模型返回的内容

        // 从响应中提取token的使用数，累加到tokenCount中
        Optional.of(chatResponse)
                .map(ChatResponse::getMetadata)//提取元数据
                .map(ChatResponseMetadata::getUsage)//提取使用信息包含 Token 使用情况
                .map(Usage::getTotalTokens)//提取总 Token 数
                .ifPresent(tokenCount::addAndGet);//累加 Token 数 如果前面的操作都成功（即所有对象都不为 null），则执行 tokenCount.addAndGet() 方法

        // 提取出大模型的响应内容，进行后续的处理
        Optional.of(chatResponse)
                .map(ChatResponse::getResult)
                .map(Generation::getOutput)//提取输出消息
                .ifPresent(assistantMessage -> processResponse(assistantMessage, inputList));
        //如果前面的操作都成功，则执行 lambda 表达式，调用 processResponse 方法
        //将提取出的 assistantMessage 和原始的 inputList 作为参数传入
    }
    /**
     * 处理大模型返回的响应内容
     * 处理过程：
     *
     * 检测到工具调用请求
     * 确认工具名称为"AgentOutput"
     * 解析参数为ReActInput对象
     * 调用agentOutput方法执行搜索操作
     * 通过这种方式，processResponse方法有效地将大模型的抽象响应转换为代理的具体行动，是实现ReAct代理模式的核心环节。
     */
    private void processResponse(AssistantMessage assistantMessage, List<Message> inputList) {
        var text = assistantMessage.getText();//模型返回的助手消息对象，包含文本内容和可能的工具调用请求
        log.info("开始处理大模型的响应，内容：{}", text);
        if (assistantMessage.hasToolCalls()) {//首先检查响应中是否包含工具调用请求
            // 如果有多个工具调用请求，记录警告日志（因为当前实现主要处理单个工具调用）
            if (assistantMessage.getToolCalls().size() > 1) {
                log.warn("找到1个以上的工具执行请求列表:{}", assistantMessage.getToolCalls());
            }
            // 获取第一个工具调用请求
            var tool = CollUtil.getFirst(assistantMessage.getToolCalls());
            // 检查工具名称是否为AGENT_OUTPUT_METHOD（常量，值为"AgentOutput"）
            if (StrUtil.equalsIgnoreCase(tool.name(), AGENT_OUTPUT_METHOD)) {
                // 如果匹配，将工具参数转换为ReActInput对象，并调用agentOutput方法执行具体操作
                var reActInput = JSONUtil.toBean(tool.arguments(), ReActInput.class);
                this.agentOutput(reActInput, inputList, assistantMessage);
            } else {
                //如果工具名称不匹配，记录错误日志
                log.error("未知工具名称:{}", tool.name());
            }
            return;
        }

        // 如果没有工具调用，则尝试从响应中解析出ReActInput
        //方法从文本中提取第一个JSON格式的字符串
        JsonFinder.findFirstJson(text)
                .map(json -> JSONUtil.toBean(json, ReActInput.class))
                .ifPresent(reActInput -> this.agentOutput(reActInput, inputList, assistantMessage));
    }

    /*
    实际工作场景举例
    假设代理正在执行一个搜索任务，代码执行过程如下：

    场景：用户要求代理搜索"2023年Java发展趋势"

    大模型生成动作：{"search": {"query": "2023年Java发展趋势", "maxResults": 3}}
    代码提取动作名称：actName = "search"
    从actionRegistry获取search函数
    准备参数：{"query": "2023年Java发展趋势", "maxResults": 3}
    执行搜索函数，获取结果："2023年Java的主要发展趋势包括：1. Java 21的发布..."
    记录结果：{"search": {...}}->2023年Java的主要发展趋势包括...
    这个结果会被添加到对话历史中，供大模型下一步思考使用
    通过这种机制，代理能够将大模型的抽象思考转化为具体的行动，并将行动结果反馈给大模型，形成完整的思考-行动-观察循环
    */
    private void agentOutput(ReActInput reActInput, List<Message> inputList, AssistantMessage assistantMessage) {
        var status = reActInput.getStatus();//从大模型响应中解析出的结构化数据，包含代理的思考过程和要执行的动作
        log.info("🏹 previous goal:{}", status.getEvaluationPreviousGoal()); //前一个目标的评估结果
        log.info("💾 memory:{}", status.getMemory());//代理的记忆内容
        log.info("🧠 thinking:{}", status.getThinking());//当前的思考过程

        // 向客户端发送消息  将思考过程发送给客户端，使客户端能够实时了解代理的工作状态。
        this.messageSession.sendMessage(String.format("[%s]%s", this.name(), status.getThinking()));

        // 获取动作列表
        var srcActionList = reActInput.getAction(); //首先获取动作列表并记录日志
        log.info("🛠 action:{}", reActInput.getAction());

        //创建两个容器：processedActionList用于跟踪已执行的动作，
        // resultBuilder用于构建执行结果
        var processedActionList = new LinkedList<Map<String, Object>>();
        var resultBuilder = new StringBuilder();

        // 1.run all action
        for (var action : srcActionList) {
            processedActionList.add(action); //将动作添加到已处理列表
            //提取动作名称（假设每个动作Map只有一个键）
            /*action.keySet()获取这个Map中所有的键（理论上只有一个）
            .stream().findFirst()将键集合转换为流并获取第一个元素（即动作名称）
            .orElseThrow()如果没有找到任何键，则抛出异常（确保动作名称必须存在）
             举个例子：如果action的结构是{"search": {"query": "Java ReAct模式"}}，那么actName就是"search"。
             */
            var actName = action.keySet().stream().findFirst().orElseThrow();
            try {
                //从actionRegistry中获取对应的执行函数并执行
                /*
                  它存储了动作名称到执行函数的映射
                  actionRegistry.get(actName)根据动作名称获取对应的执行函数
                  BeanUtil.beanToMap()将参数对象转换为Map<String, Object>类型
                  apply调用执行函数，并传入转换后的参数Map，执行函数处理后返回一个String类型的结果
                  这个结果被赋值给actionResult变量
                 */
                var actionResult = actionRegistry.get(actName).apply(BeanUtil.beanToMap(action.get(actName)));
                //将动作和结果追加到结果构建器
                resultBuilder.append(action).append("->").append(actionResult).append('\n');
            } catch (Exception e) {
                log.error("error in running action:{}", action, e);
                // break loop
                //异常处理：如执行出错，记录错误并中断循环
                resultBuilder.append(action).append("->").append("Exception:").append(StrUtil.truncateUtf8(e.getMessage(), 32)).append('\n');
                break;
            }
            if (srcActionList.indexOf(action) != srcActionList.size() - 1 && this.isStatusSignificantChanged()) {
                // 如果当前动作不是最后一个动作，并且状态发生了显著变化，则通过 break 提前终止循环。
                break;
            }
        }

        if (processedActionList.size() < srcActionList.size()) {
            // 如果已经执行的动作数量小于总动作数量，则说明状态发生了显著变化，需要忽略已经执行的动作。
            // 注意：前面的操作会导致上下文发生重大变化，忽略以下操作：
            resultBuilder.append("NOTICE:previous action leads a significant change on context,ignore following action:");
            var ignore = new ArrayList<>(srcActionList);
            ignore.removeAll(processedActionList);
            ignore.forEach(act -> {
                resultBuilder.append(act);
                resultBuilder.append('\n');
            });
            // 您应该根据最新状态继续操作
            resultBuilder.append("You should continue you step based on the newest status");
        }

        // 2.得到结果
        var actionResult = resultBuilder.toString();
        log.info("📃 obs:{}", actionResult);

        // 3.按当前结果重建消息列表（历史部分）
        /*
         这部分代码的核心作用是重构对话历史，为ReAct代理的下一轮思考-行动循环做准备。
         具体来说，它将代理执行动作的结果整合到对话历史中，形成完整的思考-行动-观察闭环。
         */
        // 删除最新的聊天消息（UserMessage）
        inputList.remove(inputList.size() - 1);
        // 把响应加入到输入列表中
        inputList.add(assistantMessage);
        if (assistantMessage.hasToolCalls()) {
            var toolCall = CollUtil.getFirst(assistantMessage.getToolCalls());
            // 添加工具消息响应
            var toolResponseMessage = new ToolResponseMessage(
                    List.of(new ToolResponseMessage.ToolResponse(toolCall.id(), toolCall.name(), "success")));
            inputList.add(toolResponseMessage);
        }
        inputList.add(new UserMessage(String.format("思路'%s'的行动结果如下:\n%s", status.getThinking(), actionResult)));
    }

    /**
     * 这是一个保护方法，默认返回false，表示状态未发生显著变化。
     * 子类可以重写此方法，根据具体业务场景实现自定义的状态变化检测逻辑
     * @return
     */
    protected boolean isStatusSignificantChanged() {
        return false;
    }

    private void buildCurrentParse(List<Message> inputList) {
        inputList.add(new UserMessage(this.getCurrentStatus()));
    }

    private List<ToolCallback> getToolCallbacks() {
        return List.of(new ToolCallback() {
            //返回当前代理中定义的toolDefinition对象
            //这个对象包含了工具的元数据信息（如工具名称、描述、参数结构等）
            //主要用于向大模型描述可用的工具能力，使大模型了解可以调用哪些工具
            @Override
            public @NotNull ToolDefinition getToolDefinition() {
                return toolDefinition;
            }

            //当大模型决定调用工具时，理论上会触发此方法执行具体工具逻辑
            //此实现中主动抛出了UnsupportedOperationException异常
            /*
             * 异常信息"toolCalling should controlled by native program"表明：
             * 工具调用流程不应通过此回调机制执行
             * 工具调用逻辑应由原生程序控制（可能在代理的其他部分实现）
             */
            @Override
            public @NotNull String call(@NotNull String toolInput) {
                throw new UnsupportedOperationException("toolCalling should controlled by native program.");
            }
        });
    }
    /**
     * 获取任务当前的状态数据，使用L2M规则每次放入上下文的最末端
     *
     * @return
     */
    protected abstract String getCurrentStatus();

    /**
     * 1.增加SystemMessage
     * 2.增加主任务对应的UserMessage
     * 3.增加One-Shot Example（可选）
     */
    protected abstract void addInitInput(List<Message> inputList, String task);

    /**
     * 初始化action方法绑定，应与构造参数中传入的schema匹配
     */
    private void initActionMap() {
        /**
         * 场景：当 ReAct 代理完成任务并需要结束思考-行动-观察循环时
         * 触发条件：大模型认为任务已完成，决定调用 "done" 工具来结束代理运行
         * 参数传递：
         * 大模型生成包含最终答案的参数，格式为：{"text": "这是最终的答案内容"}
         * 这个 JSON 对象被转换为 Map<String, Object> 类型的 paramMap
         * // paramMap 的实际内容
         * paramMap = {"text": "这是最终的答案内容"}
         *
         * // 1. 从 paramMap 中获取 text 参数
         * var result = MapUtil.getStr(paramMap, "text"); // result = "这是最终的答案内容"
         *
         * // 2. 调用 done 方法设置最终结果并停止代理
         * this.done(result); // 内部会设置 finalAnswer = result 并设置 isRunning = false
         *
         * // 3. 返回结果供调用方使用
         * return result; // 返回 "这是最终的答案内容"
         */
        actionRegistry.put("done", paramMap -> {
            var result = MapUtil.getStr(paramMap, "text");
            this.done(result);
            return result;
        });
        toolCallbackProvider().forEach(this::mergeByToolCallbackProvider);
    }
    protected List<ToolCallbackProvider> toolCallbackProvider() {
        return List.of();
    }
    /**
     * 将已有Tool的定义并入ReAct的Schema中，例如MCP服务提供的能力
     */
    private void mergeByToolCallbackProvider(ToolCallbackProvider toolCallback) {
        // 读取原始的工具定义
        JSONObject jsonObject = JSONUtil.parseObj(this.toolSchema);
        String expression = "inputSchema.properties.action.items.properties";
        // 读取具体工具的定义
        JSONObject actionEntry = jsonObject.getByPath(expression, JSONObject.class);
        for (ToolCallback tool : toolCallback.getToolCallbacks()) {
            var def = tool.getToolDefinition(); //获取每个工具的定义
            // 将子类中定义的工具加入到工具定义中
            actionEntry.set(def.name(), JSONUtil.parseObj(def.inputSchema()));//将工具的输入模式（inputSchema）添加到 JSON 对象中
            actionRegistry.put(def.name(), mp -> tool.call(JSONUtil.toJsonStr(mp)));
            /**
             * 具体工作场景举例
             * 假设代理有一个名为 search 的搜索工具，其输入模式要求 query（搜索关键词）和 max_results（最大结果数）两个参数：
             * 大模型决定使用 search 工具，并生成符合要求的参数
             * 这些参数被封装成一个 Map 对象，例如：{"query": "人工智能", "max_results": 5}
             * 这个 Map 对象就是 mp，被传递给注册的 Lambda 表达式
             * Lambda 表达式将 mp 转换为 JSON 字符串："{"query": "人工智能", "max_results": 5}"
             * 最后调用 tool.call("{"query": "人工智能", "max_results": 5}") 执行实际的搜索操作
             */
        }
        // 回写到原始对象中
        jsonObject.putByPath(expression, actionEntry);
        this.toolSchema = jsonObject.toString();
        this.refreshToolDefinition();
    }

    /**
     * 完成动作
     */
    protected void done(String result) {
        this.agentStop = true;
        this.finalResult = result;
    }
    @Override
    public ChatModel chatModel() {
        return this.chatModel;
    }
    public void resetReActConfig(Consumer<ReActConfig> cons) {
        cons.accept(reActConfig);
    }
}
